import base64
import io
import time
import uuid
from pathlib import Path
from typing import Optional, List

import aiofiles
import aiohttp
from PIL import Image
from fastapi import APIRouter
from fastapi import File, UploadFile, Form, HTTPException
from loguru import logger
from paddleocr import PaddleOCR
from pydantic import BaseModel
from pydantic import HttpUrl

api_router = APIRouter(tags=["main"])

UPLOAD_DIR: Path = Path(__file__).parent.parent.joinpath("tmp")
UPLOAD_DIR.mkdir(exist_ok=True)


class OCRHandler:
    """OCR处理器，用于管理全局OCR实例"""

    _instance = None
    _ocr = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(OCRHandler, cls).__new__(cls)
        return cls._instance

    def get_ocr(self):
        """
        获取OCR实例，延迟初始化

        Returns:
            PaddleOCR: OCR实例
        """
        if self._ocr is None:
            self._ocr = PaddleOCR(
                # 关闭文档整体自动旋转
                use_doc_orientation_classify=False,
                # 关闭文件自动去畸变
                use_doc_unwarping=False,
                # 关闭文本行自动旋转
                use_textline_orientation=False,
                # 测试环境初次部署时注释以下配置, 测试好后, 生产将~/.paddlex下model复制到自定义目录, 可离线部署
                # text_detection_model_dir="D:/Home/PaddleOCR_Models/PP-OCRv5_server_det",
                # text_recognition_model_dir="D:/Home/PaddleOCR_Models/PP-OCRv5_server_rec",
            )
        return self._ocr


# 创建全局实例
ocr_handler = OCRHandler()


# 为了向后兼容，提供一个函数接口
def get_ocr():
    """获取OCR实例的函数接口"""
    global ocr_handler
    return ocr_handler.get_ocr()


# 直接提供全局ocr_instance实例
ocr_instance = ocr_handler.get_ocr()


class PaddleOcrResult(BaseModel):
    text_list: List[str]
    poly_list: List[List[List[float]]]


def get_ocr_result(image_path: Path) -> PaddleOcrResult | None:
    start_time = time.time()
    result = ocr_instance.predict(str(image_path))
    end_time = time.time()
    recognition_time = end_time - start_time

    if not result:
        return None

    res = result[0]
    text_list = res.get("rec_texts")
    poly_list = res.get("rec_polys")
    # 转换numpy数组为可序列化的列表
    if poly_list is not None:
        poly_list = [poly.tolist() for poly in poly_list]
    logger.info(f"{text_list=}, {poly_list=}")
    logger.info(f"OCR识别结果：{text_list}，识别耗时：{recognition_time:.2f}秒")
    image_path.unlink()
    return PaddleOcrResult(text_list=text_list, poly_list=poly_list)


class OCRResponse(BaseModel):
    success: bool
    result: Optional[PaddleOcrResult] = None
    error: Optional[str] = None


async def save_uploaded_file(file: UploadFile) -> Path:
    """保存上传的文件"""
    filename = f"{uuid.uuid4()}_{file.filename}"
    file_path = UPLOAD_DIR.joinpath(filename)
    async with aiofiles.open(str(file_path), "wb") as out_file:
        content = await file.read()
        await out_file.write(content)
    return file_path


@api_router.post("/file", response_model=OCRResponse)
async def ocr_file_endpoint(file: UploadFile = File(...)):
    try:
        start_time = time.time()
        image_path = await save_uploaded_file(file)
        ocr_result = get_ocr_result(image_path)
        cost = time.time() - start_time
        logger.info(f"处理完成，耗时：{cost:.2f}秒")
        return OCRResponse(success=True, result=ocr_result)
    except Exception as e:
        logger.exception(f"处理失败: {str(e)}")
        return OCRResponse(success=False, error=f"处理失败: {str(e)}")


async def save_image_from_url(url: HttpUrl) -> Path:
    """从URL下载并保存图片"""
    async with aiohttp.ClientSession() as session:
        async with session.get(str(url)) as response:
            if response.status != 200:
                raise HTTPException(status_code=400, detail="无法从URL获取图片")

            content = await response.read()
            image = Image.open(io.BytesIO(content))
            filename = f"{uuid.uuid4()}_url_image.jpg"
            file_path = UPLOAD_DIR.joinpath(filename)
            image.save(file_path)
            return file_path


@api_router.post("/url", response_model=OCRResponse)
async def ocr_url_endpoint(url: HttpUrl = Form(...)):
    try:
        start_time = time.time()
        image_path = await save_image_from_url(url)
        ocr_result = get_ocr_result(image_path)
        cost = time.time() - start_time
        logger.info(f"处理完成，耗时：{cost:.2f}秒")
        return OCRResponse(success=True, result=ocr_result)
    except Exception as e:
        logger.exception(f"处理失败: {str(e)}")
        return OCRResponse(success=False, error=f"处理失败: {str(e)}")


def save_image_from_base64(base64_str: str) -> Path:
    """从Base64字符串保存图片"""
    try:
        # 处理可能存在的data URI前缀
        if "," in base64_str:
            _, base64_str = base64_str.split(",", 1)

        image_data = base64.b64decode(base64_str)
        image = Image.open(io.BytesIO(image_data))
        filename = f"{uuid.uuid4()}_base64_image.jpg"
        file_path = UPLOAD_DIR.joinpath(filename)
        image.save(file_path)
        return file_path
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Base64解码失败: {str(e)}")


@api_router.post("/base64", response_model=OCRResponse)
async def ocr_base64_endpoint(base64_image: str = Form(...)):
    try:
        start_time = time.time()
        image_path = save_image_from_base64(base64_image)
        ocr_result = get_ocr_result(image_path)
        cost = time.time() - start_time
        logger.info(f"处理完成，耗时：{cost:.2f}秒")
        return OCRResponse(success=True, result=ocr_result)
    except Exception as e:
        logger.exception(f"处理失败: {str(e)}")
        return OCRResponse(success=False, error=f"处理失败: {str(e)}")
